# coding: utf-8
import os
import json
import time
import math
import matplotlib.pyplot as plt
import numpy as np
from numpy import newaxis

from core.data_processor import DataLoader
from core.model import Model
from keras.models import load_model


def plot_results(predicted_data, true_data):
    plt.plot(predicted_data, color='r')
    # plt.plot(true_data)
    plt.show()


def plot_results_multiple(predicted_data, true_data, prediction_len):
    fig = plt.figure(facecolor='white')
    ax = fig.add_subplot(111)
    ax.plot(true_data, label='True Data')
    # Pad the list of predictions to shift it in the graph to it's correct start
    for i, data in enumerate(predicted_data):
        padding = [None for p in range(i * prediction_len)]
        plt.plot(padding + data, label='Prediction')
        plt.legend()
    plt.show()


def predict_sequences_multiple(model, data, window_size, prediction_len):
    print('[Model] Predicting Sequences Multiple...')
    prediction_seqs = []
    for i in range(int(len(data) / prediction_len)):
        curr_frame = data[i * prediction_len]
        predicted = []
        for j in range(prediction_len):
            predicted.append(model.predict(curr_frame[newaxis, :, :])[0, 0])
            curr_frame = curr_frame[1:]
            curr_frame = np.insert(curr_frame, [window_size - 2], predicted[-1], axis=0)
        prediction_seqs.append(predicted)
    return prediction_seqs


def main():
    model_name = '17082021-110312-e10.h5'
    configs = json.load(open('config.json', 'r'))
    if not os.path.exists(configs['model']['save_dir']):
        os.makedirs(configs['model']['save_dir'])

    data = DataLoader(
        os.path.join('data', configs['data']['filename']),
        configs['data']['train_test_split'],
        configs['data']['columns']
    )
    model = load_model(os.path.join(configs['model']['save_dir'], model_name))

    x_test, y_test = data.get_test_data(
        seq_len=configs['data']['sequence_length'],
        normalise=configs['data']['normalise']
    )
    predictions = predict_sequences_multiple(model, x_test, configs['data']['sequence_length'],
                                             configs['data']['sequence_length'])
    plot_results(predictions, y_test)


if __name__ == '__main__':
    main()
